#! /usr/bin/python
# encoding=utf-8
# Created by Fenglu Niu on 2025/4/30 11:07
import asyncio
import functools
import json
import os

from dotenv import load_dotenv
from fengluU import exec_time
from langchain_core.messages import SystemMessage, HumanMessage
from langchain_core.prompts import PromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_mcp_adapters.client import MultiServerMCPClient
from langchain_openai import ChatOpenAI
from langgraph.prebuilt import create_react_agent

import common


class LLM(object):

    def __init__(self, temperature=0.3):
        load_dotenv()
        base_url = os.getenv('OPENAI_API_BASE')
        model_deepseek = os.getenv('OPENAI_MODEL_DEEPSEEK')
        model_qwen = os.getenv('OPENAI_MODEL_QWEN')
        self.__deepseek = ChatOpenAI(
            base_url=base_url,
            model=model_deepseek,
            temperature=temperature
        )
        self.__qwen = ChatOpenAI(
            base_url=base_url,
            model=model_qwen,
            temperature=temperature
        )

    @property
    def deepseek(self):
        return self.__deepseek

    @property
    def qwen(self):
        return self.__qwen

    def dispatch_agent(self, text):
        prompt = PromptTemplate.from_file(f'{common.root_path()}/prompt/dispatch_agent_prompt.md')
        messages = [
            SystemMessage(content=prompt.template),
            HumanMessage(content=text)
        ]
        response = self.__deepseek.invoke(messages)
        content = response.content.strip().replace('<think>.*?</think>\n\n', '')
        return json.loads(content)

    def gen_title(self, query):
        prompt = PromptTemplate.from_file(f'{common.root_path()}/prompt/gen_title_prompt.md')
        messages = [
            SystemMessage(content=prompt.template),
            HumanMessage(content=query)
        ]
        response = self.__deepseek.invoke(messages)
        content = response.content.strip().replace('<think>\n\n</think>\n\n', '')
        return content

if __name__ == '__main__':
    llm = LLM()
    print(llm.dispatch_agent('我现在心情很差，除了日常上课外，还可以做那些事情充实大学生活？'))
    # print(llm.gen_title('我现在心情很差，我应该做什么？'))
    # asyncio.run(llm.stest("我现在心情很差，我应该做什么？"))
